Skip to content

Commit f0d65b6

Browse files
cjpaiscjpais
authored
Merge pull request #717 from corebonts/gemma3
Initial support for Gemma 3 models
2 parents 17d7f4a + 4d20599 commit f0d65b6

File tree

1 file changed

+231
-1
lines changed

1 file changed

+231
-1
lines changed

‎llama.cpp/llama.cpp

Lines changed: 231 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ enum llm_arch {
162162
LLM_ARCH_MINICPM,
163163
LLM_ARCH_GEMMA,
164164
LLM_ARCH_GEMMA2,
165+
LLM_ARCH_GEMMA3,
165166
LLM_ARCH_STARCODER2,
166167
LLM_ARCH_MAMBA,
167168
LLM_ARCH_XVERSE,
@@ -211,6 +212,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
211212
{ LLM_ARCH_MINICPM, "minicpm" },
212213
{ LLM_ARCH_GEMMA, "gemma" },
213214
{ LLM_ARCH_GEMMA2, "gemma2" },
215+
{ LLM_ARCH_GEMMA3, "gemma3" },
214216
{ LLM_ARCH_STARCODER2, "starcoder2" },
215217
{ LLM_ARCH_MAMBA, "mamba" },
216218
{ LLM_ARCH_XVERSE, "xverse" },
@@ -1008,6 +1010,26 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
10081010
{ LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" },
10091011
},
10101012
},
1013+
{
1014+
LLM_ARCH_GEMMA3,
1015+
{
1016+
{ LLM_TENSOR_TOKEN_EMBD, "token_embd" },
1017+
{ LLM_TENSOR_OUTPUT_NORM, "output_norm" },
1018+
{ LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" },
1019+
{ LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" },
1020+
{ LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" },
1021+
{ LLM_TENSOR_ATTN_K, "blk.%d.attn_k" },
1022+
{ LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" },
1023+
{ LLM_TENSOR_ATTN_V, "blk.%d.attn_v" },
1024+
{ LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" },
1025+
{ LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" },
1026+
{ LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" },
1027+
{ LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" },
1028+
{ LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" },
1029+
{ LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" },
1030+
{ LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" },
1031+
},
1032+
},
10111033
{
10121034
LLM_ARCH_STARCODER2,
10131035
{
@@ -1935,6 +1957,7 @@ struct llama_hparams {
19351957
uint32_t n_layer;
19361958
uint32_t n_rot;
19371959
uint32_t n_swa = 0; // sliding window attention (SWA)
1960+
uint32_t n_swa_pattern = 1; // by default, all layers use non-sliding-window attention
19381961
uint32_t n_embd_head_k; // dimension of keys (d_k). d_q is assumed to be the same, but there are n_head q heads, and only n_head_kv k-v heads
19391962
uint32_t n_embd_head_v; // dimension of values (d_v) aka n_embd_head
19401963
uint32_t n_expert = 0;
@@ -1962,7 +1985,9 @@ struct llama_hparams {
19621985

19631986
float rope_attn_factor = 1.0f;
19641987
float rope_freq_base_train;
1988+
float rope_freq_base_train_swa;
19651989
float rope_freq_scale_train;
1990+
float rope_freq_scale_train_swa;
19661991
uint32_t n_ctx_orig_yarn;
19671992
float rope_yarn_log_mul;
19681993

@@ -2001,6 +2026,7 @@ struct llama_hparams {
20012026
if (this->n_layer != other.n_layer) return true;
20022027
if (this->n_rot != other.n_rot) return true;
20032028
if (this->n_swa != other.n_swa) return true;
2029+
if (this->n_swa_pattern != other.n_swa_pattern) return true;
20042030
if (this->n_embd_head_k != other.n_embd_head_k) return true;
20052031
if (this->n_embd_head_v != other.n_embd_head_v) return true;
20062032
if (this->n_expert != other.n_expert) return true;
@@ -2035,6 +2061,8 @@ struct llama_hparams {
20352061
if (!is_float_close(this->rope_attn_factor, other.rope_attn_factor, EPSILON)) return true;
20362062
if (!is_float_close(this->rope_freq_base_train, other.rope_freq_base_train, EPSILON)) return true;
20372063
if (!is_float_close(this->rope_freq_scale_train, other.rope_freq_scale_train, EPSILON)) return true;
2064+
if (!is_float_close(this->rope_freq_base_train_swa, other.rope_freq_base_train_swa, EPSILON)) return true;
2065+
if (!is_float_close(this->rope_freq_scale_train_swa, other.rope_freq_scale_train_swa, EPSILON)) return true;
20382066
if (!is_float_close(this->expert_weights_scale, other.expert_weights_scale, EPSILON)) return true;
20392067
if (!is_float_close(this->rope_yarn_log_mul, other.rope_yarn_log_mul, EPSILON)) return true;
20402068
if (!is_float_close(this->f_residual_scale, other.f_residual_scale, EPSILON)) return true;
@@ -4446,6 +4474,10 @@ static void llm_load_hparams(
44464474
}
44474475
hparams.rope_freq_scale_train = ropescale == 0.0f ? 1.0f : 1.0f/ropescale;
44484476

4477+
// by default assume that the sliding-window layers use the same scaling type as the non-sliding-window layers
4478+
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
4479+
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
4480+
44494481
ml.get_key(LLM_KV_ROPE_SCALING_ATTN_FACTOR, hparams.rope_attn_factor, false);
44504482

44514483
// non-transformer models do not have attention heads
@@ -4779,6 +4811,8 @@ static void llm_load_hparams(
47794811
case LLM_ARCH_GEMMA2:
47804812
{
47814813
hparams.n_swa = 4096; // default value of gemma 2
4814+
hparams.n_swa_pattern = 2;
4815+
47824816
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
47834817
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
47844818
ml.get_key(LLM_KV_ATTN_LOGIT_SOFTCAPPING, hparams.f_attn_logit_softcapping, false);
@@ -4792,6 +4826,28 @@ static void llm_load_hparams(
47924826
default: model.type = e_model::MODEL_UNKNOWN;
47934827
}
47944828
} break;
4829+
case LLM_ARCH_GEMMA3:
4830+
{
4831+
hparams.n_swa_pattern = 6;
4832+
4833+
hparams.rope_freq_base_train_swa = 10000.0f;
4834+
hparams.rope_freq_scale_train_swa = 1.0f;
4835+
4836+
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
4837+
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
4838+
4839+
switch (hparams.n_layer) {
4840+
case 26: model.type = e_model::MODEL_1B; break;
4841+
case 34: model.type = e_model::MODEL_4B; break;
4842+
case 48: model.type = e_model::MODEL_12B; break;
4843+
case 62: model.type = e_model::MODEL_27B; break;
4844+
default: model.type = e_model::MODEL_UNKNOWN;
4845+
}
4846+
4847+
hparams.f_attention_scale = model.type == e_model::MODEL_27B
4848+
? 1.0f / std::sqrt(float(hparams.n_embd / hparams.n_head(0)))
4849+
: 1.0f / std::sqrt(float(hparams.n_embd_head_k));
4850+
} break;
47954851
case LLM_ARCH_STARCODER2:
47964852
{
47974853
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
@@ -6844,6 +6900,38 @@ static bool llm_load_tensors(
68446900
layer.ffn_post_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd});
68456901
}
68466902
} break;
6903+
case LLM_ARCH_GEMMA3:
6904+
{
6905+
model.tok_embd = ml.create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
6906+
6907+
// output
6908+
model.output_norm = ml.create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
6909+
model.output = ml.create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); // same as tok_embd, duplicated to allow offloading
6910+
6911+
for (int i = 0; i < n_layer; ++i) {
6912+
ggml_context * ctx_layer = ctx_for_layer(i);
6913+
ggml_context * ctx_split = ctx_for_layer_split(i);
6914+
6915+
auto & layer = model.layers[i];
6916+
6917+
layer.attn_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
6918+
6919+
layer.wq = ml.create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
6920+
layer.wk = ml.create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
6921+
layer.wv = ml.create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
6922+
layer.wo = ml.create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
6923+
6924+
layer.attn_post_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);
6925+
layer.attn_k_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
6926+
layer.attn_q_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
6927+
6928+
layer.ffn_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
6929+
layer.ffn_gate = ml.create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
6930+
layer.ffn_up = ml.create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
6931+
layer.ffn_down = ml.create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
6932+
layer.ffn_post_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);
6933+
}
6934+
} break;
68476935
case LLM_ARCH_STARCODER2:
68486936
{
68496937
model.tok_embd = ml.create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab});
@@ -11727,7 +11815,8 @@ struct llm_build_context {
1172711815

1172811816
for (int il = 0; il < n_layer; ++il) {
1172911817
// (il % 2) layers use SWA
11730-
struct ggml_tensor * KQ_mask_l = (il % 2 == 0) ? KQ_mask_swa : KQ_mask;
11818+
const bool is_swa = il % hparams.n_swa_pattern < (hparams.n_swa_pattern - 1);
11819+
struct ggml_tensor * KQ_mask_l = is_swa ? KQ_mask_swa : KQ_mask;
1173111820

1173211821
// norm
1173311822
cur = llm_build_norm(ctx0, inpL, hparams,
@@ -11839,6 +11928,142 @@ struct llm_build_context {
1183911928
return gf;
1184011929
}
1184111930

11931+
struct ggml_cgraph * build_gemma3() {
11932+
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false);
11933+
11934+
const int64_t n_embd_head_k = hparams.n_embd_head_k;
11935+
11936+
struct ggml_tensor * cur;
11937+
struct ggml_tensor * inpL;
11938+
11939+
inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
11940+
11941+
// TODO: do not normalize weights for raw embeddings input (i.e. encoded image emdeddings)
11942+
inpL = ggml_scale(ctx0, inpL, sqrtf(n_embd));
11943+
cb(inpL, "inp_scaled", -1);
11944+
11945+
// inp_pos - contains the positions
11946+
struct ggml_tensor * inp_pos = build_inp_pos();
11947+
11948+
struct ggml_tensor * KQ_mask = build_inp_KQ_mask(true);
11949+
struct ggml_tensor * KQ_mask_swa = build_inp_KQ_mask_swa(true);
11950+
11951+
for (int il = 0; il < n_layer; ++il) {
11952+
const bool is_swa = il % hparams.n_swa_pattern < (hparams.n_swa_pattern - 1);
11953+
11954+
const float freq_base_l = is_swa ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base;
11955+
const float freq_scale_l = is_swa ? hparams.rope_freq_scale_train_swa : cparams.rope_freq_scale;
11956+
11957+
struct ggml_tensor * KQ_mask_l = is_swa ? KQ_mask_swa : KQ_mask;
11958+
11959+
// norm
11960+
cur = llm_build_norm(ctx0, inpL, hparams,
11961+
model.layers[il].attn_norm, NULL,
11962+
LLM_NORM_RMS, cb, il);
11963+
cb(cur, "attn_norm", il);
11964+
11965+
// self-attention
11966+
{
11967+
// compute Q and K and RoPE them
11968+
struct ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur);
11969+
cb(Qcur, "Qcur", il);
11970+
11971+
struct ggml_tensor * Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur);
11972+
cb(Kcur, "Kcur", il);
11973+
11974+
struct ggml_tensor * Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, cur);
11975+
cb(Vcur, "Vcur", il);
11976+
11977+
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head, n_tokens);
11978+
Qcur = llm_build_norm(ctx0, Qcur, hparams,
11979+
model.layers[il].attn_q_norm, NULL,
11980+
LLM_NORM_RMS, cb, il);
11981+
cb(Qcur, "Qcur_normed", il);
11982+
11983+
Qcur = ggml_rope_ext(
11984+
ctx0, Qcur, inp_pos, nullptr,
11985+
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
11986+
ext_factor, attn_factor, beta_fast, beta_slow);
11987+
cb(Qcur, "Qcur", il);
11988+
11989+
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv, n_tokens);
11990+
Kcur = llm_build_norm(ctx0, Kcur, hparams,
11991+
model.layers[il].attn_k_norm, NULL,
11992+
LLM_NORM_RMS, cb, il);
11993+
cb(Kcur, "Kcur_normed", il);
11994+
11995+
Kcur = ggml_rope_ext(
11996+
ctx0, Kcur, inp_pos, nullptr,
11997+
n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
11998+
ext_factor, attn_factor, beta_fast, beta_slow);
11999+
cb(Kcur, "Kcur", il);
12000+
12001+
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
12002+
model.layers[il].wo, NULL,
12003+
Kcur, Vcur, Qcur, KQ_mask_l, n_tokens, kv_head, n_kv, hparams.f_attention_scale, cb, il);
12004+
}
12005+
12006+
cur = llm_build_norm(ctx0, cur, hparams,
12007+
model.layers[il].attn_post_norm, NULL,
12008+
LLM_NORM_RMS, cb, il);
12009+
cb(cur, "attn_post_norm", il);
12010+
12011+
if (il == n_layer - 1) {
12012+
// skip computing output for unused tokens
12013+
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
12014+
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
12015+
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
12016+
}
12017+
12018+
struct ggml_tensor * sa_out = ggml_add(ctx0, cur, inpL);
12019+
cb(sa_out, "sa_out", il);
12020+
12021+
cur = llm_build_norm(ctx0, sa_out, hparams,
12022+
model.layers[il].ffn_norm, NULL,
12023+
LLM_NORM_RMS, cb, il);
12024+
cb(cur, "ffn_norm", il);
12025+
12026+
// feed-forward network
12027+
{
12028+
cur = llm_build_ffn(ctx0, lctx, cur,
12029+
model.layers[il].ffn_up, NULL, NULL,
12030+
model.layers[il].ffn_gate, NULL, NULL,
12031+
model.layers[il].ffn_down, NULL, NULL,
12032+
NULL,
12033+
LLM_FFN_GELU, LLM_FFN_PAR, cb, il);
12034+
cb(cur, "ffn_out", il);
12035+
}
12036+
12037+
cur = llm_build_norm(ctx0, cur, hparams,
12038+
model.layers[il].ffn_post_norm, NULL,
12039+
LLM_NORM_RMS, cb, -1);
12040+
cb(cur, "ffn_post_norm", -1);
12041+
12042+
cur = ggml_add(ctx0, cur, sa_out);
12043+
cur = lctx.cvec.apply_to(ctx0, cur, il);
12044+
cb(cur, "l_out", il);
12045+
12046+
// input for next layer
12047+
inpL = cur;
12048+
}
12049+
12050+
cur = inpL;
12051+
12052+
cur = llm_build_norm(ctx0, cur, hparams,
12053+
model.output_norm, NULL,
12054+
LLM_NORM_RMS, cb, -1);
12055+
12056+
cb(cur, "result_norm", -1);
12057+
12058+
// lm_head
12059+
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
12060+
12061+
cb(cur, "result_output", -1);
12062+
12063+
ggml_build_forward_expand(gf, cur);
12064+
12065+
return gf;
12066+
}
1184212067

1184312068
struct ggml_cgraph * build_starcoder2() {
1184412069
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false);
@@ -14136,6 +14361,10 @@ static struct ggml_cgraph * llama_build_graph(
1413614361
{
1413714362
result = llm.build_gemma2();
1413814363
} break;
14364+
case LLM_ARCH_GEMMA3:
14365+
{
14366+
result = llm.build_gemma3();
14367+
} break;
1413914368
case LLM_ARCH_STARCODER2:
1414014369
{
1414114370
result = llm.build_starcoder2();
@@ -17315,6 +17544,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) {
1731517544
case LLM_ARCH_PHI3:
1731617545
case LLM_ARCH_GEMMA:
1731717546
case LLM_ARCH_GEMMA2:
17547+
case LLM_ARCH_GEMMA3:
1731817548
case LLM_ARCH_STARCODER2:
1731917549
case LLM_ARCH_OPENELM:
1732017550
case LLM_ARCH_GPTNEOX:

0 commit comments

Comments
 (0)